/*
 * HSM Proxy Project.
 * Copyright (C) 2013 FedICT.
 * Copyright (C) 2013 Frank Cornelis.
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License version
 * 3.0 as published by the Free Software Foundation.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, see 
 * http://www.gnu.org/licenses/.
 */

package be.fedict.hsm.client;

import java.io.ByteArrayInputStream;
import java.net.MalformedURLException;
import java.net.ProxySelector;
import java.security.PrivateKey;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;

import javax.xml.bind.JAXBElement;
import javax.xml.ws.Binding;
import javax.xml.ws.BindingProvider;
import javax.xml.ws.handler.Handler;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import be.fedict.hsm.ws.DSSConstants;
import be.fedict.hsm.ws.DigitalSignatureServiceFactory;
import be.fedict.hsm.ws.ResultMajor;
import be.fedict.hsm.ws.jaxb.dss.AnyType;
import be.fedict.hsm.ws.jaxb.dss.Base64Signature;
import be.fedict.hsm.ws.jaxb.dss.DocumentHash;
import be.fedict.hsm.ws.jaxb.dss.InputDocuments;
import be.fedict.hsm.ws.jaxb.dss.KeySelector;
import be.fedict.hsm.ws.jaxb.dss.ObjectFactory;
import be.fedict.hsm.ws.jaxb.dss.ResponseBaseType;
import be.fedict.hsm.ws.jaxb.dss.Result;
import be.fedict.hsm.ws.jaxb.dss.SignRequest;
import be.fedict.hsm.ws.jaxb.dss.SignResponse;
import be.fedict.hsm.ws.jaxb.dss.SignatureObject;
import be.fedict.hsm.ws.jaxb.hsm.GetAliasesRequest;
import be.fedict.hsm.ws.jaxb.hsm.GetCertificateChainRequest;
import be.fedict.hsm.ws.jaxb.xmldsig.DigestMethodType;
import be.fedict.hsm.ws.jaxb.xmldsig.KeyInfoType;
import be.fedict.hsm.ws.jaxb.xmldsig.X509DataType;
import be.fedict.hsm.ws.jaxws.DigitalSignatureService;
import be.fedict.hsm.ws.jaxws.DigitalSignatureServicePortType;

/**
 * The HSM Proxy web service client.
 * 
 * @author Frank Cornelis
 * 
 */
public class HSMProxyClient {

	private static final Log LOG = LogFactory.getLog(HSMProxyClient.class);

	private final static Map<String, String> digestAlgoToDigestMethod;

	static {
		digestAlgoToDigestMethod = new HashMap<String, String>();
		digestAlgoToDigestMethod.put("SHA1",
				"http://www.w3.org/2000/09/xmldsig#sha1");
		digestAlgoToDigestMethod.put("SHA-1",
				"http://www.w3.org/2000/09/xmldsig#sha1");
		digestAlgoToDigestMethod.put("SHA-256",
				"http://www.w3.org/2001/04/xmlenc#sha256");
		digestAlgoToDigestMethod.put("SHA-512",
				"http://www.w3.org/2001/04/xmlenc#sha512");
	}

	private final DigitalSignatureServicePortType dssPort;

	private final ObjectFactory dssObjectFactory;

	private final be.fedict.hsm.ws.jaxb.xmldsig.ObjectFactory dsObjectFactory;

	private final be.fedict.hsm.ws.jaxb.hsm.ObjectFactory hsmObjectFactory;

	private final CertificateFactory certificateFactory;

	private static final ClientProxySelector clientProxySelector;

	private final String endpointAddress;

	static {
		ProxySelector defaultProxySelector = ProxySelector.getDefault();
		clientProxySelector = new ClientProxySelector(defaultProxySelector);
		ProxySelector.setDefault(clientProxySelector);
	}

	/**
	 * Main constructor. To access the HSM Proxy web service, you need to have a
	 * valid credential.
	 * 
	 * @param endpointAddress
	 *            the HSM Proxy web service endpoint address.
	 * @param credentialPrivateKey
	 *            the credential private key.
	 * @param credentialCertificate
	 *            the corresponding credential X509 certificate.
	 */
	public HSMProxyClient(String endpointAddress,
			PrivateKey credentialPrivateKey,
			X509Certificate credentialCertificate) {
		this.endpointAddress = endpointAddress;

		DigitalSignatureService digitalSignatureService = DigitalSignatureServiceFactory
				.getInstance();
		this.dssPort = digitalSignatureService.getDigitalSignatureServicePort();

		BindingProvider bindingProvider = (BindingProvider) this.dssPort;
		bindingProvider.getRequestContext().put(
				BindingProvider.ENDPOINT_ADDRESS_PROPERTY, endpointAddress);

		this.dssObjectFactory = new ObjectFactory();
		this.dsObjectFactory = new be.fedict.hsm.ws.jaxb.xmldsig.ObjectFactory();
		this.hsmObjectFactory = new be.fedict.hsm.ws.jaxb.hsm.ObjectFactory();

		try {
			this.certificateFactory = CertificateFactory.getInstance("X.509");
		} catch (CertificateException e) {
			throw new RuntimeException("X509 certificate factory error: "
					+ e.getMessage(), e);
		}

		Binding binding = bindingProvider.getBinding();
		List<Handler> handlerChain = binding.getHandlerChain();
		handlerChain.add(new WSSecuritySOAPHandler(credentialPrivateKey,
				credentialCertificate));
		binding.setHandlerChain(handlerChain);
	}

	/**
	 * Sets the HTTP proxy that should be used to communicate with the HSP Proxy
	 * web service.
	 * 
	 * @param proxyHost
	 *            the HTTP proxy host.
	 * @param proxyPort
	 *            the HTTP proxy port number.
	 */
	public void setProxy(String proxyHost, int proxyPort) {
		try {
			clientProxySelector.setProxy(this.endpointAddress, proxyHost,
					proxyPort);
		} catch (MalformedURLException e) {
			LOG.error("URL error: " + e.getMessage(), e);
		}
	}

	/**
	 * Signs the given digest value via the HSM Proxy web service.
	 * 
	 * @param digestValue
	 *            the digest value.
	 * @param digestAlgo
	 *            the digest algorithm. For example "SHA-1".
	 * @param keyAlias
	 *            the key alias that has been registered within the HSM Proxy.
	 * @return
	 */
	public byte[] sign(byte[] digestValue, String digestAlgo, String keyAlias) {
		SignRequest signRequest = this.dssObjectFactory.createSignRequest();
		signRequest.setProfile(DSSConstants.HSM_PROXY_DSS_PROFILE_URI);
		String requestId = "request-" + UUID.randomUUID().toString();
		signRequest.setRequestID(requestId);
		InputDocuments inputDocuments = this.dssObjectFactory
				.createInputDocuments();
		signRequest.setInputDocuments(inputDocuments);

		DocumentHash documentHash = this.dssObjectFactory.createDocumentHash();
		inputDocuments.getDocumentOrTransformedDataOrDocumentHash().add(
				documentHash);

		DigestMethodType dsDigestMethod = this.dsObjectFactory
				.createDigestMethodType();
		documentHash.setDigestMethod(dsDigestMethod);
		String digestMethod = digestAlgoToDigestMethod.get(digestAlgo);
		if (null == digestMethod) {
			throw new IllegalArgumentException("unsupported digest algo: "
					+ digestAlgo);
		}
		dsDigestMethod.setAlgorithm(digestMethod);

		documentHash.setDigestValue(digestValue);

		AnyType optionalInputs = this.dssObjectFactory.createAnyType();
		signRequest.setOptionalInputs(optionalInputs);

		KeySelector keySelector = this.dssObjectFactory.createKeySelector();
		optionalInputs.getAny().add(keySelector);
		KeyInfoType keyInfo = this.dsObjectFactory.createKeyInfoType();
		keySelector.setKeyInfo(keyInfo);
		keyInfo.getContent().add(this.dsObjectFactory.createKeyName(keyAlias));

		optionalInputs.getAny().add(
				this.dssObjectFactory.createSignatureType("urn:ietf:rfc:3447"));

		SignResponse signResponse = this.dssPort.sign(signRequest);

		Result result = signResponse.getResult();
		String resultMajor = result.getResultMajor();
		if (false == ResultMajor.SUCCESS.getUri().equals(resultMajor)) {
			throw new RuntimeException("error occurred");
		}

		SignatureObject signatureObject = signResponse.getSignatureObject();
		Base64Signature base64Signature = signatureObject.getBase64Signature();
		byte[] signatureValue = base64Signature.getValue();
		return signatureValue;
	}

	/**
	 * Gives back the list of key aliases available within the HSM Proxy key
	 * store for the set application credential.
	 * 
	 * @return the list of aliases.
	 */
	public Set<String> getAliases() {
		GetAliasesRequest getAliasesRequest = this.hsmObjectFactory
				.createGetAliasesRequest();
		getAliasesRequest.setProfile(DSSConstants.HSM_PROXY_DSS_PROFILE_URI);
		String requestId = "request-" + UUID.randomUUID().toString();
		getAliasesRequest.setRequestID(requestId);
		ResponseBaseType response = this.dssPort.getAliases(getAliasesRequest);
		// TODO: error handling
		List<Object> optionalOutputContentList = response.getOptionalOutputs()
				.getAny();
		Set<String> aliases = new HashSet<String>();
		for (Object optionalOutputContent : optionalOutputContentList) {
			LOG.debug("optional output type: "
					+ optionalOutputContent.getClass().getName());
			if (optionalOutputContent instanceof KeySelector) {
				KeySelector keySelector = (KeySelector) optionalOutputContent;
				KeyInfoType keyInfo = keySelector.getKeyInfo();
				List<Object> keyInfoContent = keyInfo.getContent();
				for (Object keyInfoObject : keyInfoContent) {
					if (keyInfoObject instanceof JAXBElement) {
						JAXBElement jaxbElement = (JAXBElement) keyInfoObject;
						String alias = (String) jaxbElement.getValue();
						aliases.add(alias);
					}
				}
			}
		}
		return aliases;
	}

	/**
	 * Gives back the X509 certificate chain for the given key alias.
	 * 
	 * @param alias
	 * @return the X509 certificate chain as a list.
	 * @throws CertificateException
	 */
	public List<X509Certificate> getCertificateChain(String alias)
			throws CertificateException {
		GetCertificateChainRequest getCertificateChainRequest = this.hsmObjectFactory
				.createGetCertificateChainRequest();
		getCertificateChainRequest
				.setProfile(DSSConstants.HSM_PROXY_DSS_PROFILE_URI);
		String requestId = "request-" + UUID.randomUUID().toString();
		getCertificateChainRequest.setRequestID(requestId);

		{
			AnyType optionalInputs = this.dssObjectFactory.createAnyType();
			getCertificateChainRequest.setOptionalInputs(optionalInputs);
			KeySelector keySelector = this.dssObjectFactory.createKeySelector();
			optionalInputs.getAny().add(keySelector);
			KeyInfoType keyInfo = this.dsObjectFactory.createKeyInfoType();
			keySelector.setKeyInfo(keyInfo);
			keyInfo.getContent().add(this.dsObjectFactory.createKeyName(alias));
		}

		ResponseBaseType response = this.dssPort
				.getCertificateChain(getCertificateChainRequest);
		// TODO: error handling
		List<Object> optionalOutputContentList = response.getOptionalOutputs()
				.getAny();
		for (Object optionalOutputContent : optionalOutputContentList) {
			LOG.debug("optional output type: "
					+ optionalOutputContent.getClass().getName());
			if (optionalOutputContent instanceof JAXBElement) {
				JAXBElement jaxbElement = (JAXBElement) optionalOutputContent;
				KeyInfoType keyInfo = (KeyInfoType) jaxbElement.getValue();
				List<Object> keyInfoContent = keyInfo.getContent();
				JAXBElement keyInfoObject = (JAXBElement) keyInfoContent.get(0);
				LOG.debug("key info object type: "
						+ keyInfoObject.getClass().getName());
				X509DataType x509Data = (X509DataType) keyInfoObject.getValue();
				List<Object> x509DataContent = x509Data
						.getX509IssuerSerialOrX509SKIOrX509SubjectName();
				List<X509Certificate> certificateChain = new LinkedList<X509Certificate>();
				for (Object x509DataObject : x509DataContent) {
					LOG.debug("x509 data object type: "
							+ x509DataObject.getClass().getName());
					JAXBElement x509DataElement = (JAXBElement) x509DataObject;
					LOG.debug("type: "
							+ x509DataElement.getValue().getClass().getName());
					byte[] encodedCertificate = (byte[]) x509DataElement
							.getValue();
					X509Certificate certificate = (X509Certificate) this.certificateFactory
							.generateCertificate(new ByteArrayInputStream(
									encodedCertificate));
					certificateChain.add(certificate);
				}
				return certificateChain;
			}
		}
		return null;
	}
}
